# tests/test_ml_fill.py
import pytest
import pandas as pd
from ml_fill import predict_missing_with_rf

@pytest.fixture
def sample_data():
    return pd.DataFrame({
        '国名Ch': ['中国', '美国', '日本', '中国', '美国'],
        'Year': [2020, 2020, 2020, 2021, 2021],
        'GDP': [14.7, 21.4, 5.0, 15.2, None],
        '人口': [1400, 328, 126, 1412, 330]
    })

def test_predict_missing_with_rf(sample_data):
    # 测试是否能成功预测而不报错
    filled_df = predict_missing_with_rf(sample_data.copy(), 'GDP')
    
    # 检查缺失值是否被填补
    assert filled_df['GDP'].isna().sum() == 0
    
    # 检查原始非缺失值未被修改
    assert filled_df.loc[0, 'GDP'] == 14.7
    assert filled_df.loc[1, 'GDP'] == 21.4